#!/bin/bash
category="GPU"
cycle=${1:-'60'}
randomNum=${2:-'default'}
log="/tmp/gpu-burn.log"
plog="/tmp/gpu-power.log"
err="/tmp/gpu-burn.err"
runlog="/tmp/gpu-run.log"

gpuCount=${3:-'8'}
minGflops=${4:-'0'}
maxTemp=${5:-'0'}
minPower=${6:-'0'}
arpAvg=${7:-'0'}

resultcode=0
resultmsg=""   # Perf: GE ${minGflops} Gflops

start_time=$(date +%s)

# GPU数量   > $runlog
gpu_count=$(nvidia-smi --list-gpus | wc -l)
echo "GPU压测开始，GPU数量：" $gpu_count > $runlog

if [[ $gpuCount -ne $gpu_count ]]; then
	resultcode="6"
	resultmsg="GPU quantity error, actual ${gpu_count}, configured ${gpuCount}"

	end_time=$(date +%s)

	run_time=$((end_time - start_time))

	echo "category=${category};resultcode=${resultcode};resultmsg=${resultmsg};runtime=${run_time}" >> $runlog
	echo "category=${category};resultcode=${resultcode};resultmsg=${resultmsg};runtime=${run_time}"
	exit 1
fi

# 启动 gpu_burn 命令并将输出重定向到 burn.log
cd /root/scripts/gputest/gpu-burn-master
# nvidia-smi -pm 1 >> $runlog
./gpu_burn -tc $cycle > $log 2>&1 &
gpu_burn_pid=$!  # 获取 gpu_burn 进程的 PID
echo gpu_burn_pid= $gpu_burn_pid >> $runlog


# 性能监视函数
perfor_monitor() {
	# 从接口来  还有 ./build/all_reduce_perf -b 128M -e 1G -f 2 -g 8 平均值   -g 后面的是gpu的数量
	count=0;
    tail -f $log | while IFS= read -r line; do
		if [ -f /tmp/resultcode.txt ]; then
			exit 1
		fi
		if [[ $line == *"proc'd:"* ]]; then
		
			IFS='%' # 设置分隔符为 '%'
			for entry in $line; do
				if [[ -n $entry ]]; then

					count=$((count + 1))
					if [ "$count" -le 100 ]; then
						continue  # 忽略前100行
					fi

					# 处理一行的N个GPU
					# 提取 Gflop/s 和温度
					gflops_array=($(echo "$entry" | grep -oP '\(\K[0-9]+(?= Gflop)'))

					# 提取温度
					temps_array=($(echo "$entry" | grep -oP '\d+ C' | sed 's/ C//'))

					#echo "----" $gflops_array "====" $temps_array "----"
					# 设置IFS为换行符
					echo "$gflops_array" | while IFS= read -r item; do
						if [ "$item" != "" ]; then
							if [[ $item -le $minGflops ]]; then
								echo "Gflops条件未满足，退出 gpu_burn...Gflops=$item--" >> $runlog

								echo "9" > /tmp/resultcode.txt
								echo "Pref：$item/${minGflops} Gflops" > /tmp/resultmsg.txt

								if kill -0 $gpu_burn_pid 2>/dev/null; then
									kill $gpu_burn_pid 2>> $runlog # 退出 gpu_burn 进程
								fi
								exit 1
							fi
						fi
						
					done
					echo "$temps_array" | while IFS= read -r item; do
						if [ "$item" != "" ]; then
							if [[ $item -ge $maxTemp ]]; then
								echo "Temp条件未满足，退出 gpu_burn...Temp=$item" >> $runlog

								echo "8" > /tmp/resultcode.txt
								echo "Pref：$item/${maxTemp} Temp" > /tmp/resultmsg.txt

								if kill -0 $gpu_burn_pid 2>/dev/null; then
									kill $gpu_burn_pid 2>> $runlog # 退出 gpu_burn 进程
								fi
								exit 1
							fi
						fi
						
					done

					# -lt：小于
					# -le：小于或等于
					# -gt：大于
					# -ge：大于或等于
					# -eq：等于
					# -ne：不等于

					
					# echo "temps_array 的长度为: ${#temps_array[@]}"

					# for i in "${!temps_array[@]}"; do
					# 	echo "GPU $((i+1)): 温度: ${temps_array[$i]}======="
					# done

				fi
			done


		fi
        
    done
}

# 功耗监视函数
power_monitor(){

	sleep 20
	# nvidia-smi | grep "Default" | awk '{print $5}'| sed 's/W//'
	# 记录一个开始时间 然后使用执行完时间判断是否需要退出
	current_timestamp=$(date +%s)
	exit_time=$((current_timestamp + cycle - 20))

	# 减20 是为了开始和结束各留10秒不检测
	pmtime=$((cycle - 30))
	# 清下日志
	>$plog
	for i in $(seq 1 $pmtime);
	do
		if [ -f /tmp/resultcode.txt ]; then
			exit 1
		fi
		# echo "这是第 $i 次循环"
		# 获取功耗值并去掉单位
		# 新的获取功耗命令：nvidia-smi --query-gpu=power.draw --format=csv,noheader,nounits | awk '{print int($1)}'

		# 使用nvidia-smi命令获取功耗，并将结果存储在数组中
		power_draw=($(nvidia-smi --query-gpu=power.draw --format=csv,noheader,nounits | awk '{print int($1)}'))
		# 记录功耗日志
		nvidia-smi >> $plog

		for power in "${power_draw[@]}"; do
			echo $power "==========" $i >> $runlog
		done
		
		if [ "$minPower" -gt 0 ]; then
			# 循环遍历数组，并显示功耗
			for power in "${power_draw[@]}"; do
				# echo "当前功耗: ${power} W"
				if [[ $power -le $minPower ]]; then
					echo "Power条件未满足，退出 gpu_burn...Power=$power" >> $runlog

					echo "7" > /tmp/resultcode.txt
					echo "Pref：$power/${minPower} Power" > /tmp/resultmsg.txt

					if kill -0 $gpu_burn_pid 2>/dev/null; then
						kill $gpu_burn_pid 2>> $runlog # 退出 gpu_burn 进程
					fi
					exit 1
				fi
			done
		fi

		# 等待1秒
		sleep 1

		curr_time=$(date +%s)
		if (($curr_time > $exit_time)); then
			break
		fi
	done

	echo "计时结束，退出功耗检测。" >> $runlog

	sleep 10
	# 启动功率监视函数  1秒执行一次判断功耗
	#nvidia-smi 
}

rm -rf /tmp/resultcode.txt
rm -rf /tmp/resultmsg.txt

if [ "$minGflops" -gt 0 ]; then
	# 启动性能监视函数
	# minGflops  maxTemp 存在这两个值则开启 性能检测线程
	resultmsg="Perf: GE ${minGflops} Gflops"
	perfor_monitor &
	perfor_monitor_pid=$!
	echo perfor_monitor_pid= $perfor_monitor_pid >> $runlog
fi

# 启动功耗监视函数
power_monitor &
power_monitor_pid=$!
echo power_monitor_pid= $power_monitor_pid >> $runlog

wait $power_monitor_pid 2>> $runlog

wait $gpu_burn_pid 2>> $runlog

pkill -f "tail -f $log" 2>> $runlog

# 验证gpu_burn log文件最后的GPU状态

# 提取 GPU 编号（如 0,1,2...）
gpu_ids=$(grep -o "GPU [0-9]\+" $log | awk '{print $2}' | sort -u)

# 检查每个 GPU 的状态
all_ok=true
for gpu_id in $gpu_ids; do
    if grep -q "GPU $gpu_id: OK" $log; then
        echo "GPU $gpu_id: OK" >> $runlog
    else
        echo "GPU $gpu_id: NOT OK!" >> $runlog
        all_ok=false
    fi
done

if $all_ok; then
    echo "All GPUs are OK." >> $runlog
else
    echo "Warning: Some GPUs are not OK!" >> $runlog
    echo "6" > /tmp/resultcode.txt
	echo "Pref：GPU not OK" > /tmp/resultmsg.txt
fi



# # 停止监视   没有必要了  肯定是退出的
# if kill -0 $perfor_monitor_pid 2>/dev/null; then
#     kill -9 $perfor_monitor_pid
# fi
#  
# if kill -0 $power_monitor_pid 2>/dev/null; then
#     kill -9 $power_monitor_pid
# fi

if [ -f /tmp/resultcode.txt ]; then
    resultcode=$(< /tmp/resultcode.txt)
	resultmsg=$(< /tmp/resultmsg.txt)
# else
# 	if (( $(echo "$arpAvg > 0" | bc -l) )); then
# 		sleep 30
# 		# 启动 
# 		cd /root/scripts/gputest/nccl-tests-master/build
# 		./all_reduce_perf -b 128M -e 1G -f 2 -g $gpu_count >> $runlog


# 		echo "GPU通讯能力测试，执行完成======" >> $runlog

# 		currAvg=$(grep "Avg bus bandwidth" $runlog | awk '{print $NF}')

# 		# 判断 currAvg 是否小于 currAvg
# 		if (( $(echo "$currAvg < $arpAvg" | bc -l) )); then
# 			resultcode="9"
# 			# resultmsg="检测未通过，检测值（$currAvg）,阈值（$arpAvg）"
# 			resultmsg="Pref：LE ${arpAvg} ARP"
# 		fi
# 	fi
	
fi

end_time=$(date +%s)

run_time=$((end_time - start_time))

echo "GPU测试，执行完成，执行时间：${run_time}======" >> $runlog


# echo "category=${category};resultcode=${resultcode};resultmsg=${resultmsg};runtime=${run_time}"

for ((i=1;i<=gpu_count;i++));do
    model=$(nvidia-smi -i $((i-1)) --query-gpu=name --format=csv,noheader)
    echo "category=${category};resultcode=${resultcode};resultmsg=${resultmsg};runtime=${run_time};model=${model}"
done



exit 0


